# -*- coding: utf-8 -*-

# 大模型提取关键词

import numpy as np
import jieba
from transformers import pipeline

class PromptKeyExtractor:
    def __init__(self, model='distilbert-base-uncased'):
        self.keyword_extractor = pipeline('feature-extraction', model=model)

    def extract_keywords(self, text, top_n=5):
        """提取文本的关键词"""
        # 生成特征向量
        features = self.keyword_extractor(text)
        # 获取关键词
        keywords = self._get_top_keywords(features, text, top_n)
        return keywords

    def _get_top_keywords(self, features, text, top_n):
        """从特征向量中提取关键词"""
        feature_array = np.array(features).squeeze(0)  # 取出第一个（即唯一一个）序列
        scores = feature_array.mean(axis=0)  # 计算每个词的平均得分

        # 使用结巴分词
        words = list(jieba.cut(text))  # 使用结巴分词处理中文文本
        
        # 生成 (词, 得分) 对
        word_scores = {word: scores[i] for i, word in enumerate(words) if i < len(scores)}

        # 按得分排序，提取得分最高的词
        sorted_keywords = sorted(word_scores.items(), key=lambda x: x[1], reverse=True)[:top_n]
        return [(keyword, score) for keyword, score in sorted_keywords]

# 使用示例
if __name__ == "__main__":
    text = "我爱自然语言处理和机器学习。机器学习是一种数据分析方法，它使用算法来从数据中学习。"
    
    # 创建 PromptKey 提取器实例
    extractor = PromptKeyExtractor()
    
    # 提取关键词
    keywords = extractor.extract_keywords(text)
    print("提取的关键词：", keywords)
